package pl.edu.agh.bp;

import pl.edu.agh.neural.core.INeuron;
import pl.edu.agh.neural.core.ITrainableLayerWithTeacher;
import pl.edu.agh.neural.simple.InputConnection;
import pl.edu.agh.neural.simple.SimpleLayer;
import pl.edu.agh.neural.simple.SimpleNeuron;

import java.util.List;

public class BpLayer extends SimpleLayer implements ITrainableLayerWithTeacher {

    public BpLayer(List<SimpleNeuron> neurons, boolean hasBias) {
        super(neurons, hasBias);
    }

    @Override
    public void train(double learningSpeed, double momentum, double[] expectedValues) {
        int neuronsCount = neurons.size();
        if (expectedValues.length != neuronsCount) {
            throw new RuntimeException("Expected values vector length is different than number of output neurons");
        }

        for (int i = 0; i < neuronsCount; i++) {
            INeuron neuron = neurons.get(i);
            double z = expectedValues[i];
            double y = neuron.getValue();
            //System.out.println("Expected: " + z + " Was: " + y);
            double dy = neuron.getDerivative();
            double error = dy * (z - y);

            setNewWeights(neuron, error, learningSpeed, momentum);
        }
    }

    @Override
    public void train(double learningSpeed, double momentum, ITrainableLayerWithTeacher nextLayer) {
        for (INeuron neuron : neurons) {
            double dy = neuron.getDerivative();
            double sum = 0.0;
            for (INeuron nextLayerNeuron : nextLayer.getNeurons()) {
                for (InputConnection inputConnection : nextLayerNeuron.getInputConnections()) {
                    if (inputConnection.getNeuron() == neuron) {
                        sum += inputConnection.getWeight() * nextLayerNeuron.getError();
                    }
                }
            }

            double error = dy * sum;
            setNewWeights(neuron, error, learningSpeed, momentum);
        }
    }

    private void setNewWeights(INeuron neuron, double error, double learningSpeed, double momentum) {
        neuron.setError(error);
        for (InputConnection inputConnection : neuron.getInputConnections()) {
            double x = inputConnection.getNeuron().getValue();
            double dw = inputConnection.getWeightDelta();
            double newWeight = inputConnection.getWeight() + learningSpeed * error * x + momentum * dw;
            inputConnection.setWeight(newWeight);
        }
    }
}
